// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/builder_run_end.h>
#include <nebula/array/builder_primitive.h>

#include <cstddef>
#include <cstdint>
#include <utility>
#include <vector>

#include <nebula/types/scalar.h>
#include <turbo/base/checked_cast.h>
#include <nebula/numeric/int_util_overflow.h>
#include <turbo/log/logging.h>
#include <nebula/util/ree_util.h>

namespace nebula::internal {

    RunCompressorBuilder::RunCompressorBuilder(MemoryPool *pool,
                                               std::shared_ptr<ArrayBuilder> inner_builder,
                                               std::shared_ptr<DataType> type)
            : ArrayBuilder(pool), inner_builder_(std::move(inner_builder)) {}

    RunCompressorBuilder::~RunCompressorBuilder() = default;

    void RunCompressorBuilder::reset() {
        current_run_length_ = 0;
        current_value_.reset();
        inner_builder_->reset();
        UpdateDimensions();
    }

    turbo::Status RunCompressorBuilder::ResizePhysical(int64_t capacity) {
        TURBO_RETURN_NOT_OK(inner_builder_->resize(capacity));
        UpdateDimensions();
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::append_nulls(int64_t length) {
        if (TURBO_UNLIKELY(length == 0)) {
            return turbo::OkStatus();
        }
        if (TURBO_UNLIKELY(current_run_length_ == 0)) {
            // open a new NULL run
                    DKCHECK_EQ(current_value_, nullptr);
            current_run_length_ = length;
        } else if (current_value_ == nullptr) {
            // Extend the currently open NULL run
            current_run_length_ += length;
        } else {
            // close then non-NULL run
            TURBO_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
            TURBO_RETURN_NOT_OK(inner_builder_->append_scalar(*current_value_));
            UpdateDimensions();
            // open a new NULL run
            current_value_.reset();
            current_run_length_ = length;
        }
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::append_empty_values(int64_t length) {
        if (TURBO_UNLIKELY(length == 0)) {
            return turbo::OkStatus();
        }
        // Empty values are usually appended as placeholders for future values, so
        // we make no attempt at making the empty values appended now part of the
        // current run. Each append_empty_values() creates its own run of the given length.
        TURBO_RETURN_NOT_OK(FinishCurrentRun());
        {
            TURBO_RETURN_NOT_OK(WillCloseRunOfEmptyValues(length));
            TURBO_RETURN_NOT_OK(inner_builder_->append_empty_value());
            UpdateDimensions();
        }
        // Current run remains cleared after FinishCurrentRun() as we don't want to
        // extend it with empty values potentially coming in the future.
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::append_scalar(const Scalar &scalar, int64_t n_repeats) {
        if (TURBO_UNLIKELY(n_repeats == 0)) {
            return turbo::OkStatus();
        }
        if (TURBO_UNLIKELY(current_run_length_ == 0)) {
            // open a new run
            current_value_ = scalar.is_valid ? scalar.shared_from_this() : nullptr;
            current_run_length_ = n_repeats;
        } else if ((current_value_ == nullptr && !scalar.is_valid) ||
                   (current_value_ != nullptr && current_value_->equals(scalar))) {
            // Extend the currently open run
            current_run_length_ += n_repeats;
        } else {
            // close the current run
            TURBO_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
            TURBO_RETURN_NOT_OK(current_value_ ? inner_builder_->append_scalar(*current_value_)
                                               : inner_builder_->append_null());
            UpdateDimensions();
            // open a new run
            current_value_ = scalar.is_valid ? scalar.shared_from_this() : nullptr;
            current_run_length_ = n_repeats;
        }
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::append_scalars(const ScalarVector &scalars) {
        if (scalars.empty()) {
            return turbo::OkStatus();
        }
        TURBO_RETURN_NOT_OK(ArrayBuilder::append_scalars(scalars));
        UpdateDimensions();
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::AppendRunCompressedArraySlice(
            const ArraySpan &run_compressed_array, int64_t offset, int64_t length) {
                DKCHECK(!has_open_run());
        TURBO_RETURN_NOT_OK(inner_builder_->append_array_slice(run_compressed_array, offset, length));
        UpdateDimensions();
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::FinishCurrentRun() {
        if (current_run_length_ > 0) {
            // close the current run
            TURBO_RETURN_NOT_OK(WillCloseRun(current_value_, current_run_length_));
            TURBO_RETURN_NOT_OK(current_value_ ? inner_builder_->append_scalar(*current_value_)
                                               : inner_builder_->append_null());
            UpdateDimensions();
            // Clear the current run
            current_value_.reset();
            current_run_length_ = 0;
        }
        return turbo::OkStatus();
    }

    turbo::Status RunCompressorBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        TURBO_RETURN_NOT_OK(FinishCurrentRun());
        return inner_builder_->finish_internal(out);
    }

}  // namespace nebula::internal
namespace nebula {
    // ----------------------------------------------------------------------
    // RunEndEncodedBuilder

    RunEndEncodedBuilder::ValueRunBuilder::ValueRunBuilder(
            MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &value_builder,
            const std::shared_ptr<DataType> &value_type, RunEndEncodedBuilder &ree_builder)
            : RunCompressorBuilder(pool, value_builder, value_type), ree_builder_(ree_builder) {}

    RunEndEncodedBuilder::RunEndEncodedBuilder(
            MemoryPool *pool, const std::shared_ptr<ArrayBuilder> &run_end_builder,
            const std::shared_ptr<ArrayBuilder> &value_builder, std::shared_ptr<DataType> type)
            : ArrayBuilder(pool), type_(turbo::checked_pointer_cast<RunEndEncodedType>(type)) {
        auto value_run_builder =
                std::make_shared<ValueRunBuilder>(pool, value_builder, type_->get_value_type(), *this);
        value_run_builder_ = value_run_builder.get();
        children_ = {run_end_builder, std::move(value_run_builder)};
        UpdateDimensions(0, 0);
        null_count_ = 0;
    }

    turbo::Status RunEndEncodedBuilder::ResizePhysical(int64_t capacity) {
        TURBO_RETURN_NOT_OK(value_run_builder_->ResizePhysical(capacity));
        TURBO_RETURN_NOT_OK(run_end_builder().resize(capacity));
        UpdateDimensions(committed_logical_length_, 0);
        return turbo::OkStatus();
    }

    void RunEndEncodedBuilder::reset() {
        value_run_builder_->reset();
        run_end_builder().reset();
        UpdateDimensions(0, 0);
    }

    turbo::Status RunEndEncodedBuilder::append_nulls(int64_t length) {
        TURBO_RETURN_NOT_OK(value_run_builder_->append_nulls(length));
        UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::append_empty_values(int64_t length) {
        TURBO_RETURN_NOT_OK(value_run_builder_->append_empty_values(length));
                DKCHECK_EQ(value_run_builder_->open_run_length(), 0);
        UpdateDimensions(committed_logical_length_, 0);
        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::append_scalar(const Scalar &scalar, int64_t n_repeats) {
        if (scalar.type->id() == Type::RUN_END_ENCODED) {
            return append_scalar(*turbo::checked_cast<const RunEndEncodedScalar &>(scalar).value,
                                n_repeats);
        }
        TURBO_RETURN_NOT_OK(value_run_builder_->append_scalar(scalar, n_repeats));
        UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::append_scalars(const ScalarVector &scalars) {
        TURBO_RETURN_NOT_OK(this->ArrayBuilder::append_scalars(scalars));
        UpdateDimensions(committed_logical_length_, value_run_builder_->open_run_length());
        return turbo::OkStatus();
    }

    template<typename RunEndCType>
    turbo::Status RunEndEncodedBuilder::DoAppendArraySlice(const ArraySpan &array, int64_t offset,
                                                           int64_t length) {
                DKCHECK(offset + length <= array.length);
                DKCHECK_GT(length, 0);
                DKCHECK(!value_run_builder_->has_open_run());

        ree_util::RunEndEncodedArraySpan<RunEndCType> ree_span(array, array.offset + offset,
                                                               length);
        const int64_t physical_offset = ree_span.PhysicalIndex(0);
        const int64_t physical_length =
                ree_span.PhysicalIndex(ree_span.length() - 1) + 1 - physical_offset;

        TURBO_RETURN_NOT_OK(ReservePhysical(physical_length));

        // append all the run ends from array sliced by offset and length
        for (auto it = ree_span.iterator(0, physical_offset); !it.is_end(ree_span); ++it) {
            const int64_t run_end = committed_logical_length_ + it.run_length();
            TURBO_RETURN_NOT_OK(DoAppendRunEnd<RunEndCType>(run_end));
            UpdateDimensions(run_end, 0);
        }

        // append all the values directly
        TURBO_RETURN_NOT_OK(value_run_builder_->AppendRunCompressedArraySlice(
                ree_util::ValuesArray(array), physical_offset, physical_length));

        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::append_array_slice(const ArraySpan &array, int64_t offset,
                                                         int64_t length) {
                DKCHECK(array.type->equals(type_));

        // Ensure any open run is closed before appending the array slice.
        TURBO_RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());

        if (length == 0) {
            return turbo::OkStatus();
        }

        switch (type_->run_end_type()->id()) {
            case Type::INT16:
                TURBO_RETURN_NOT_OK(DoAppendArraySlice<int16_t>(array, offset, length));
                break;
            case Type::INT32:
                TURBO_RETURN_NOT_OK(DoAppendArraySlice<int32_t>(array, offset, length));
                break;
            case Type::INT64:
                TURBO_RETURN_NOT_OK(DoAppendArraySlice<int64_t>(array, offset, length));
                break;
            default:
                return turbo::invalid_argument_error("Invalid type for run ends array: ", type_->run_end_type());
        }

        return turbo::OkStatus();
    }

    std::shared_ptr<DataType> RunEndEncodedBuilder::type() const { return type_; }

    turbo::Status RunEndEncodedBuilder::finish_internal(std::shared_ptr<ArrayData> *out) {
        // finish the values array before so we can close the current run and append
        // the last run-end.
        std::shared_ptr<ArrayData> values_data;
        TURBO_RETURN_NOT_OK(value_run_builder_->finish_internal(&values_data));
        auto values_array = make_array(values_data);

        TURBO_MOVE_OR_RAISE(auto run_ends_array, run_end_builder().finish());

        TURBO_MOVE_OR_RAISE(auto ree_array,
                            RunEndEncodedArray::create(length_, run_ends_array, values_array));
        *out = std::move(ree_array->data());
        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::FinishCurrentRun() {
        TURBO_RETURN_NOT_OK(value_run_builder_->FinishCurrentRun());
        UpdateDimensions(length_, 0);
        return turbo::OkStatus();
    }

    template<typename RunEndCType>
    turbo::Status RunEndEncodedBuilder::DoAppendRunEnd(int64_t run_end) {
        constexpr auto max = std::numeric_limits<RunEndCType>::max();
        if (TURBO_UNLIKELY(run_end > max)) {
            return turbo::invalid_argument_error("Run end value must fit on run ends type but ", run_end, " > ",
                                                 max, ".");
        }
        return turbo::checked_cast<typename CTypeTraits<RunEndCType>::BuilderType *>(
                children_[0].get())
                ->append(static_cast<RunEndCType>(run_end));
    }

    turbo::Status RunEndEncodedBuilder::AppendRunEnd(int64_t run_end) {
        switch (type_->run_end_type()->id()) {
            case Type::INT16:
                TURBO_RETURN_NOT_OK(DoAppendRunEnd<int16_t>(run_end));
                break;
            case Type::INT32:
                TURBO_RETURN_NOT_OK(DoAppendRunEnd<int32_t>(run_end));
                break;
            case Type::INT64:
                TURBO_RETURN_NOT_OK(DoAppendRunEnd<int64_t>(run_end));
                break;
            default:
                return turbo::invalid_argument_error("Invalid type for run ends array: ", type_->run_end_type());
        }
        return turbo::OkStatus();
    }

    turbo::Status RunEndEncodedBuilder::CloseRun(int64_t run_length) {
        // TODO(felipecrv): gracefully fragment runs bigger than INT32_MAX
        if (TURBO_UNLIKELY(run_length > std::numeric_limits<int32_t>::max())) {
            return turbo::invalid_argument_error(
                    "Run-length of run-encoded arrays must fit in a 32-bit signed integer.");
        }
        int64_t run_end;
        if (internal::AddWithOverflow(committed_logical_length_, run_length, &run_end)) {
            return turbo::invalid_argument_error("Run end value must fit on run ends type.");
        }
        TURBO_RETURN_NOT_OK(AppendRunEnd(/*run_end=*/run_end));
        UpdateDimensions(run_end, 0);
        return turbo::OkStatus();
    }

    ArrayBuilder &RunEndEncodedBuilder::run_end_builder() { return *children_[0]; }

    ArrayBuilder &RunEndEncodedBuilder::value_builder() { return *children_[1]; }

}  // namespace nebula
